import os.path
import cv2
import pandas as pd
from tqdm import tqdm
import json



def read_csv(csv_dir = '../data/2train_rname.csv'):
    def conver(data,save_dir):
        anno_csv = pd.DataFrame(columns=['image_name', 'x1', 'y1', 'x2', 'y2', 'class_name'])
        for index in tqdm(range(len(data))):
            name = '../../data/2_images/2_images/' + data.iloc[index,4].split('/')[-1]
            annos = json.loads(data.iloc[index,5])['items']
            for anno in annos:
                # for k,v in anno.items():
                    # print(k,'--',v)
                box = anno['meta']['geometry']
                label = str(anno['labels']).replace("{'标签': '","").replace("'}","").replace("'不合规原因'","").replace("'标签': ","")
                label = label if '{' not in label else 'wrongdressed'
                anno_csv = anno_csv.append(pd.DataFrame({'image_name':name,'x1':int(box[0]),'y1':int(box[1]),'x2':int(box[2]),'y2':int(box[3]),'class_name':label},index=[0]))
        anno_csv[['image_name','x1','y1','x2','y2','class_name']].to_csv(save_dir,index=False,header=False)

    data = pd.read_csv(csv_dir)
    data.sample(frac=1).reset_index(drop=True)
    train_data = data.iloc[:int(0.7*len(data)),:]
    val_data = data.iloc[int(0.7 * len(data)):, :]
    conver(train_data,save_dir='label_file/train.csv')
    conver(val_data,save_dir='label_file/val.csv')



def show_box(image_path,csv_dir,show_path=None):
    csv_data = pd.read_csv(csv_dir)
    last_name= csv_data.iloc[0].loc['image_name'].split('/')[-1]
    image = cv2.imread(os.path.join(image_path,last_name))
    for index in range(len(csv_data)):
        if csv_data.iloc[index].loc['image_name'] != last_name:
            print(index,last_name)
            # cv2.imshow('image with boxes',cv2.resize(image,(512,512)))
            # cv2.waitKey()
            cv2.imwrite(os.path.join(show_path,last_name),image)
            image = cv2.imread(os.path.join(image_path,csv_data.iloc[index].loc['image_name']))
            last_name = csv_data.iloc[index].loc['image_name']
        box = [int(csv_data.iloc[index, i]) for i in range(1, 5)] + [csv_data.iloc[index,5]]
        cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 5)
        cv2.putText(image,box[4],(box[0],box[1]),cv2.FONT_HERSHEY_PLAIN,10,(0,0,255),5)


if __name__ == '__main__':
    read_csv()
    # show_box(image_path='../data/2_images/2_images',csv_dir='label_file/val.csv',show_path='../data/train_show')